import { spawn } from 'child_process';
import { EventEmitter } from 'events';

class MCPServer extends EventEmitter {
  constructor(config) {
    super();
    this.name = config.name;
    this.command = config.command;
    this.args = config.args || [];
    this.env = config.env || {};
    this.description = config.description || '';
    this.process = null;
    this.connected = false;
    this.tools = [];
    this.resources = [];
    this.messageId = 1;
    this.pendingRequests = new Map();
  }

  async connect() {
    return new Promise((resolve, reject) => {
      try {
        // 启动 MCP 服务器进程
        this.process = spawn(this.command, this.args, {
          env: { ...process.env, ...this.env },
          stdio: ['pipe', 'pipe', 'pipe']
        });

        let buffer = '';
        
        this.process.stdout.on('data', (data) => {
          buffer += data.toString();
          const lines = buffer.split('\n');
          buffer = lines.pop(); // 保留不完整的行
          
          lines.forEach(line => {
            if (line.trim()) {
              this.handleMessage(line.trim());
            }
          });
        });

        this.process.stderr.on('data', (data) => {
          console.error(`[${this.name}] stderr:`, data.toString());
        });

        this.process.on('error', (error) => {
          console.error(`[${this.name}] 进程错误:`, error);
          this.connected = false;
          reject(error);
        });

        this.process.on('exit', (code) => {
          console.log(`[${this.name}] 进程退出，代码: ${code}`);
          this.connected = false;
        });

        // 发送初始化消息
        setTimeout(async () => {
          try {
            await this.initialize();
            this.connected = true;
            resolve();
          } catch (error) {
            reject(error);
          }
        }, 1000);

      } catch (error) {
        reject(error);
      }
    });
  }

  async initialize() {
    // 发送 initialize 请求
    const initResponse = await this.sendRequest('initialize', {
      protocolVersion: '2024-11-05',
      capabilities: {
        roots: { listChanged: true },
        sampling: {}
      },
      clientInfo: {
        name: 'mcp-client',
        version: '1.0.0'
      }
    });

    // 发送 initialized 通知
    this.sendNotification('notifications/initialized', {});

    // 获取工具列表
    try {
      const toolsResponse = await this.sendRequest('tools/list', {});
      this.tools = toolsResponse.tools || [];
    } catch (error) {
      console.warn(`[${this.name}] 获取工具列表失败:`, error.message);
    }

    // 获取资源列表
    try {
      const resourcesResponse = await this.sendRequest('resources/list', {});
      this.resources = resourcesResponse.resources || [];
    } catch (error) {
      console.warn(`[${this.name}] 获取资源列表失败:`, error.message);
    }
  }

  handleMessage(line) {
    try {
      const message = JSON.parse(line);
      
      if (message.id && this.pendingRequests.has(message.id)) {
        const { resolve, reject } = this.pendingRequests.get(message.id);
        this.pendingRequests.delete(message.id);
        
        if (message.error) {
          reject(new Error(message.error.message || 'MCP 服务器错误'));
        } else {
          resolve(message.result);
        }
      } else if (message.method) {
        // 处理通知
        this.emit('notification', message);
      }
    } catch (error) {
      console.error(`[${this.name}] 解析消息失败:`, error, line);
    }
  }

  sendRequest(method, params) {
    return new Promise((resolve, reject) => {
      const id = this.messageId++;
      const message = {
        jsonrpc: '2.0',
        id,
        method,
        params
      };

      this.pendingRequests.set(id, { resolve, reject });
      
      // 设置超时
      setTimeout(() => {
        if (this.pendingRequests.has(id)) {
          this.pendingRequests.delete(id);
          reject(new Error(`请求超时: ${method}`));
        }
      }, 30000);

      this.process.stdin.write(JSON.stringify(message) + '\n');
    });
  }

  sendNotification(method, params) {
    const message = {
      jsonrpc: '2.0',
      method,
      params
    };
    this.process.stdin.write(JSON.stringify(message) + '\n');
  }

  async callTool(name, args) {
    return await this.sendRequest('tools/call', {
      name,
      arguments: args
    });
  }

  async readResource(uri) {
    return await this.sendRequest('resources/read', { uri });
  }

  disconnect() {
    if (this.process) {
      this.process.kill();
      this.process = null;
    }
    this.connected = false;
  }

  getTools() {
    return this.tools;
  }

  getResources() {
    return this.resources;
  }

  isConnected() {
    return this.connected;
  }
}

export class MCPClient {
  constructor(serverConfigs, settings = {}) {
    this.servers = new Map();
    this.settings = {
      timeout: 30000,
      maxRetries: 3,
      autoConnect: true,
      logLevel: 'info',
      ...settings
    };

    // 初始化服务器
    serverConfigs.forEach(config => {
      const server = new MCPServer(config);
      this.servers.set(config.name, server);
    });
  }

  async connectAll() {
    const promises = Array.from(this.servers.values()).map(server => 
      this.connectServer(server)
    );
    
    const results = await Promise.allSettled(promises);
    
    // 检查连接结果
    results.forEach((result, index) => {
      const serverName = Array.from(this.servers.keys())[index];
      if (result.status === 'rejected') {
        console.warn(`服务器 ${serverName} 连接失败:`, result.reason.message);
      }
    });
  }

  async connectServer(server) {
    let retries = 0;
    while (retries < this.settings.maxRetries) {
      try {
        await server.connect();
        console.log(`✅ 服务器 ${server.name} 连接成功`);
        return;
      } catch (error) {
        retries++;
        console.warn(`⚠️ 服务器 ${server.name} 连接失败 (尝试 ${retries}/${this.settings.maxRetries}):`, error.message);
        
        if (retries < this.settings.maxRetries) {
          await new Promise(resolve => setTimeout(resolve, 2000 * retries));
        }
      }
    }
    throw new Error(`服务器 ${server.name} 连接失败，已达最大重试次数`);
  }

  async disconnectAll() {
    this.servers.forEach(server => {
      server.disconnect();
    });
  }

  async getAllTools() {
    const allTools = [];
    
    for (const [serverName, server] of this.servers) {
      if (server.isConnected()) {
        const tools = server.getTools();
        tools.forEach(tool => {
          allTools.push({
            ...tool,
            server: serverName,
            type: 'function',
            function: {
              name: tool.name,
              description: tool.description,
              parameters: tool.inputSchema
            }
          });
        });
      }
    }
    
    return allTools;
  }

  async callTool(toolName, args) {
    // 查找拥有该工具的服务器
    for (const [serverName, server] of this.servers) {
      if (server.isConnected()) {
        const tools = server.getTools();
        const tool = tools.find(t => t.name === toolName);
        
        if (tool) {
          try {
            const result = await server.callTool(toolName, args);
            return {
              server: serverName,
              tool: toolName,
              result: result
            };
          } catch (error) {
            throw new Error(`工具调用失败 [${serverName}:${toolName}]: ${error.message}`);
          }
        }
      }
    }
    
    throw new Error(`未找到工具: ${toolName}`);
  }

  async readResource(uri) {
    // 尝试从所有服务器读取资源
    for (const [serverName, server] of this.servers) {
      if (server.isConnected()) {
        try {
          const result = await server.readResource(uri);
          return result;
        } catch (error) {
          // 继续尝试下一个服务器
        }
      }
    }
    
    throw new Error(`无法读取资源: ${uri}`);
  }

  getServerStatus() {
    return Array.from(this.servers.values()).map(server => ({
      name: server.name,
      connected: server.isConnected(),
      description: server.description,
      toolCount: server.getTools().length,
      resourceCount: server.getResources().length
    }));
  }

  getServer(name) {
    return this.servers.get(name);
  }
}